import numpy as np
import torchvision.datasets as dsets
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from sklearn.neighbors import KNeighborsClassifier
import random

def visualize_knn_predictions(X_test, y_test, y_test_pred, num_images=10):
    """可视化KNN预测结果，显示原始图像和识别的数字"""
    # 将展平的图像重新整形为28x28
    images = X_test[:num_images].reshape(num_images, 28, 28)
    
    # 创建子图
    fig, axes = plt.subplots(2, 5, figsize=(15, 8))
    axes = axes.ravel()
    
    for i in range(num_images):
        # 显示图像
        axes[i].imshow(images[i], cmap='gray')
        
        # 设置标题显示真实标签和预测标签
        color = 'green' if y_test[i] == y_test_pred[i] else 'red'
        axes[i].set_title(f'True: {y_test[i]}, Pred: {y_test_pred[i]}', color=color, fontsize=12)
        axes[i].axis('off')
    
    plt.suptitle('KNN MNIST Digit Recognition Results', fontsize=16)
    plt.tight_layout()
    plt.show()

# 设置随机种子，每次运行时都会产生不同的随机数
random_seed = np.random.randint(0, 10000)
np.random.seed(random_seed)
random.seed(random_seed)
print(f"Random seed set to: {random_seed}")

batch_size = 100

mnist_train_dataset = dsets.MNIST(root="dataset/mnist", train=True, download=True)
mnist_test_dataset = dsets.MNIST(root="dataset/mnist", train=False, download=True)

train_loader = DataLoader(dataset=mnist_train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=mnist_test_dataset, batch_size=batch_size, shuffle=True)

# 设置随机选择的数据量
train_sample_size = 10000  # 从60000个训练样本中随机选择10000个
test_sample_size = 1000    # 从10000个测试样本中随机选择1000个

# 准备训练数据（随机选择）
print(f"Preparing training data (randomly selecting {train_sample_size} samples)...")
train_indices = np.random.choice(len(train_loader.dataset), train_sample_size, replace=False)
X_train = train_loader.dataset.data[train_indices].numpy()
# 需要转为numpy矩阵
X_train = X_train.reshape(X_train.shape[0], 28 * 28)

# 需要reshape之后才能放入knn分类器
y_train = train_loader.dataset.targets[train_indices].numpy()

# 准备测试数据（随机选择）
print(f"Preparing test data (randomly selecting {test_sample_size} samples)...")
test_indices = np.random.choice(len(test_loader.dataset), test_sample_size, replace=False)
X_test = test_loader.dataset.data[test_indices].numpy()
X_test = X_test.reshape(X_test.shape[0], 28 * 28)
y_test = test_loader.dataset.targets[test_indices].numpy()

# 使用sklearn的KNN分类器
print("Creating KNN classifier with sklearn...")
knn = KNeighborsClassifier(n_neighbors=5, algorithm='auto', n_jobs=-1)

print("Training KNN model...")
knn.fit(X_train, y_train)

print("Making predictions...")
y_test_pred = knn.predict(X_test)

num_test = y_test.shape[0]
num_correct = np.sum(y_test_pred == y_test)
accuracy = float(num_correct) / num_test

print('Got %d / %d correct => accuracy: %f' % (num_correct, num_test, accuracy))
print(f'Accuracy: {accuracy:.4f}')

# 可视化预测结果
visualize_knn_predictions(X_test, y_test, y_test_pred)